import os

from apitest.cmd.output import OutFileType, out_file_path
from apitest.cmd.parse_idl import ParseThrift, remove_keywords_suffix
from apitest.engine.const.enums import GenType

_INIT_FILE_MAP_ = {
    GenType.http: "",
    GenType.rpc: "import thriftpy2\n\nthriftpy2.install_import_hook()\n"
}


class GenCode(object):

    def __init__(self, psm, idl, out_dir):
        self.psm = psm
        self.idl = idl
        self.out_dir = out_dir
        self.thrift = ParseThrift(self.psm, self.idl)

    def generate_init(self, gen_type: GenType):
        st = _INIT_FILE_MAP_.get(gen_type, "")
        if st:
            out = out_file_path(self.out_dir, OutFileType.init)
            with open(out, "w") as f:
                f.write(_INIT_FILE_MAP_.get(gen_type, ""))
                f.close()

    # 清除临时文件
    def __temporary_files_clear(self):
        os.remove(self.thrift.new_idl)

    # 清除添加的关键字
    def __batch_remove(self):
        for filename in os.listdir(self.out_dir):
            file = os.path.join(self.out_dir, filename)
            if os.path.isfile(file):
                with open(file, "r+") as f:
                    data = f.read()
                    # 指针移到文件开头，并清空文件内容
                    f.seek(0)
                    f.truncate()
                    # 替换后再次写入
                    data = remove_keywords_suffix(data)
                    f.write(data)
                    f.close()

    # 文件清除
    def clear(self):
        self.__temporary_files_clear()
        self.__batch_remove()
